import time
import numpy as np
from environment import DriverEnv
from agent_qlearning import QlerningAgent

path = "models/Qtable.npy"
agent = QlerningAgent(path)
env = DriverEnv("maps/0")

# 训练模式- QLearning
print("StartTraining")
N = 1000
t0 = time.time()
for epoch in range(N):
    print("Training ,epoch:", epoch + 1, "/", N, ",time : ", (time.time() - t0), "s.")
    state = env.reset()

    done = False
    history = []
    k = 0.1 + 0.9 * epoch / N  # 非随机选择概率
    while not done:
        if np.random.rand() < k:
            action = agent.choose_action(state)
        else:
            action = env.action_space.sample()
        lastState = state.copy()
        state, reward, done, _ = env.step(action, epoch / N)
        history.append([lastState, action, reward, state])
    for i in range(len(history)):  # 倒着训练，方便传递
        agent.learn(
            history[-i - 1][0],
            history[-i - 1][1],
            history[-i - 1][2],
            history[-i - 1][3],
        )
    if (epoch + 1) % 500 == 0:
        # 保存训练结果
        agent.SaveTable(path)
    if (epoch + 1) % 50 == 0:
        env.mapGenerator()
